import komand
import base64
import re
from .schema import DownloadMalwareFileInput, DownloadMalwareFileOutput, Input, Output, Component
from ...util import malware


class DownloadMalwareFile(komand.Action):

    def __init__(self):
        super(self.__class__, self).__init__(
                name='download_malware_file',
                description=Component.DESCRIPTION,
                input=DownloadMalwareFileInput(),
                output=DownloadMalwareFileOutput())

    def run(self, params={}):
        response = malware.Malware(
            self.connection.config,
            params.get(Input.PROJECT_NAME),
            **{"sha256": params.get(Input.MALWARE_SHA256)}
        ).download(params.get(Input.COMPRESSOR), params.get(Input.PASSWORD))

        return {
            Output.RAW: {
                "filename": re.findall("filename=(.+)", response.headers['content-disposition'])[0],
                "content": base64.b64encode(response.content).decode("utf-8")
            }
        }
